use super::Matrix;
use std::ops::{Add, AddAssign, Index, IndexMut, Mul, MulAssign, Sub, SubAssign};

#[derive(PartialEq, Clone, Copy)]
pub struct MatrixView<'a> {
    data: &'a [f64],
    rows: usize,
    cols: usize,
    stride: usize,
}

impl<'a> MatrixView<'a> {
    pub fn new(data: &'a [f64], rows: usize, cols: usize, stride: usize) -> Self {
        Self { data, rows, cols, stride }
    }

    pub fn size(&self) -> (usize, usize) {
        (self.rows, self.cols)
    }

    pub fn row_size(&self) -> usize {
        self.rows
    }

    pub fn col_size(&self) -> usize {
        self.cols
    }

    pub fn to_mat(&self) -> Matrix {
        let mut m = Matrix::new(self.size());
        for row in 0..m.row_size() {
            for col in 0..m.col_size() {
                m[(row, col)] = self[(row, col)];
            }
        }
        m
    }

    pub fn get_row_view(&self, row: usize) -> MatrixView {
        MatrixView {
            data: &self.data[self.data_index(row, 0)..],
            rows: 1,
            cols: self.cols,
            stride: self.stride,
        }
    } 

    pub fn iter_rows(&self) -> impl Iterator<Item = MatrixView> {
        (0..self.rows).map(|row| {
            self.get_row_view(row)
        })
    }

    pub fn get_col_view(&self, col: usize) -> MatrixView {
        MatrixView {
            data: &self.data[self.data_index(0, col)..],
            rows: self.rows,
            cols: 1,
            stride: self.stride
        }
    }

    pub fn iter_cols(&self) -> impl Iterator<Item = MatrixView> {
        (0..self.cols).map(|col| {
            self.get_col_view(col)
        })
    }

    pub fn dot(self, rhs: Self) -> Matrix {
        assert_eq!(self.cols, rhs.rows);
        let mut res = Matrix::new((self.rows, rhs.cols));
        for row in 0..res.row_size() {
            for col in 0..res.col_size() {
                let s = (0..self.cols)
                    .map(|i| self[(row, i)] * rhs[(i, col)])
                    .sum::<f64>();
                res[(row, col)] = s;
            }
        }
        res
    }

    fn data_index(&self, row: usize, col: usize) -> usize {
        row * self.stride + col
    }
}

impl<'a> Index<(usize, usize)> for MatrixView<'a> {
    type Output = f64;
    fn index(&self, (row, col): (usize, usize)) -> &Self::Output {
        &self.data[self.data_index(row, col)]
    }
}

impl Add<MatrixView<'_>> for MatrixView<'_> {
    type Output = Matrix;
    fn add(self, rhs: MatrixView) -> Self::Output {
        assert_eq!(self.size(), rhs.size());

        let mut res = Matrix::new(self.size());
        for row in 0..self.rows {
            for col in 0..self.cols {
                res[(row, col)] = self[(row, col)] + rhs[(row, col)];
            }
        }
        res
    }
}

impl Mul<MatrixView<'_>> for MatrixView<'_> {
    type Output = Matrix;
    fn mul(self, rhs: MatrixView) -> Self::Output {
        assert_eq!(self.size(), rhs.size());

        let mut res = Matrix::new(self.size());
        for row in 0..self.rows {
            for col in 0..self.cols {
                res[(row, col)] = self[(row, col)] * rhs[(row, col)];
            }
        }
        res
    }
}

impl Mul<MatrixView<'_>> for f64 {
    type Output = Matrix;
    fn mul(self, rhs: MatrixView) -> Self::Output {
        let mut res = Matrix::new(rhs.size());
        for row in 0..res.row_size() {
            for col in 0..res.col_size() {
                res[(row, col)] = self * rhs[(row, col)];
            }
        }
        res
    }
}

impl Sub<MatrixView<'_>> for MatrixView<'_> {
    type Output = Matrix;
    fn sub(self, rhs: MatrixView) -> Self::Output {
        assert_eq!(self.size(), rhs.size());

        let mut res = Matrix::new(self.size());
        for row in 0..self.rows {
            for col in 0..self.cols {
                res[(row, col)] = self[(row, col)] - rhs[(row, col)];
            }
        }
        res
    }
}

#[derive(PartialEq)]
pub struct MatrixViewMut<'a> {
    data: &'a mut [f64],
    rows: usize,
    cols: usize,
    stride: usize,
}

impl<'a> MatrixViewMut<'a> {
    pub fn new(data: &'a mut [f64], rows: usize, cols: usize, stride: usize) -> Self {
        Self { data, rows, cols, stride }
    }

    pub fn size(&self) -> (usize, usize) {
        (self.rows, self.cols)
    }

    pub fn row_size(&self) -> usize {
        self.rows
    }

    pub fn col_size(&self) -> usize {
        self.cols
    }

    pub fn to_mat(&self) -> Matrix {
        let mut m = Matrix::new(self.size());
        for row in 0..m.row_size() {
            for col in 0..m.col_size() {
                m[(row, col)] = self[(row, col)];
            }
        }
        m
    }

    pub fn copy_from(&mut self, m: MatrixView) {
        assert_eq!(self.size(), m.size());
        for row in 0..self.rows {
            for col in 0..self.cols {
                self[(row, col)] = m[(row, col)];
            }
        }
    }

    pub fn sum_by(&mut self, m1: MatrixView, m2: MatrixView) {
        assert_eq!(m1.size(), m2.size());
        assert_eq!(self.size(), m1.size());
    
        for row in 0..self.rows {
            for col in 0..self.cols {
                self[(row, col)] = m1[(row, col)] + m2[(row, col)];
            }
        }
    }
    
    pub fn dot_by(&mut self, m1: MatrixView, m2: MatrixView) {
        assert_eq!(m1.size().1, m2.size().0);
        assert_eq!(self.size().0, m1.size().0);
        assert_eq!(self.size().1, m2.size().1);
    
        let size = m1.size().1;
    
        for row in 0..self.rows {
            for col in 0..self.cols {
                let v = (0..size)
                    .map(|i| m1[(row, i)] * m2[(i, col)])
                    .sum::<f64>();
                self[(row, col)] = v;
            }
        }
    }
    
    pub fn sum_eq(&mut self, m: MatrixView) {
        assert_eq!(self.size(), m.size());
    
        for row in 0..self.rows {
            for col in 0..self.cols {
                self[(row, col)] += m[(row, col)];
            }
        }
    }

    pub fn to_view(self) -> MatrixView<'a> {
        MatrixView { data: self.data, rows: self.rows, cols: self.cols, stride: self.stride }
    }

    fn data_index(&self, row: usize, col: usize) -> usize {
        row * self.stride + col
    }
}

impl<'a> Index<(usize, usize)> for MatrixViewMut<'a> {
    type Output = f64;
    fn index(&self, (row, col): (usize, usize)) -> &Self::Output {
        &self.data[self.data_index(row, col)]
    }
}

impl<'a> IndexMut<(usize, usize)> for MatrixViewMut<'a> {
    fn index_mut(&mut self, (row, col): (usize, usize)) -> &mut Self::Output {
        let index = self.data_index(row, col);
        &mut self.data[index]
    }
}

impl Add<MatrixViewMut<'_>> for MatrixViewMut<'_> {
    type Output = Matrix;
    fn add(self, rhs: MatrixViewMut) -> Self::Output {
        self.to_view() + rhs.to_view()
    }
}

impl Mul<MatrixViewMut<'_>> for MatrixViewMut<'_> {
    type Output = Matrix;
    fn mul(self, rhs: MatrixViewMut) -> Self::Output {
        self.to_view() * rhs.to_view()
    }
}

impl Sub<MatrixViewMut<'_>> for MatrixViewMut<'_> {
    type Output = Matrix;
    fn sub(self, rhs: MatrixViewMut) -> Self::Output {
        self.to_view() - rhs.to_view()
    }
}

impl AddAssign<MatrixView<'_>> for MatrixViewMut<'_> {
    fn add_assign(&mut self, rhs: MatrixView) {
        assert_eq!(self.size(), rhs.size());

        for row in 0..self.rows {
            for col in 0..self.cols {
                self[(row, col)] += rhs[(row, col)];
            }
        }
    }
} 

impl SubAssign<MatrixView<'_>> for MatrixViewMut<'_> {
    fn sub_assign(&mut self, rhs: MatrixView) {
        assert_eq!(self.size(), rhs.size());

        for row in 0..self.rows {
            for col in 0..self.cols {
                self[(row, col)] -= rhs[(row, col)];
            }
        }
    }
} 

impl MulAssign<MatrixView<'_>> for MatrixViewMut<'_> {
    fn mul_assign(&mut self, rhs: MatrixView) {
        assert_eq!(self.size(), rhs.size());

        for row in 0..self.rows {
            for col in 0..self.cols {
                self[(row, col)] *= rhs[(row, col)];
            }
        }
    }
} 

#[cfg(test)]
mod test {
    #[allow(unused)]
    use super::*;
    use super::super::{Matrix, mat};

    #[test]
    fn test_mat_sum() {
        let m1 = mat![[-1, -1, -1], [2, 4, 5]];
        let m2 = mat![[-3,  0, -1], [1, 9, 2]];
        let mut m3 = Matrix::new(m1.size());
        m3.view_mut().sum_by(m1.view(), m2.view());
        assert_eq!(m3, mat![[-4, -1, -2], [3, 13, 7]]);
    }

    #[test]
    fn test_mat_dot() {
        let m1 = Matrix::e(3);
        let m2 = mat![[1, 2, 3], [4, 5, 6], [7, 8, 9]];
        let mut m3 = Matrix::new(m2.size());
        m3.view_mut().dot_by(m1.view(), m2.view());
        assert_eq!(m2, m3);
    }
}